<!DOCTYPE html>
<html lang="zh-cn" dir="ltr">
    <head><meta charset='utf-8'>
<meta name='viewport' content='width=device-width, initial-scale=1'><meta name='description' content='LLama系列模型 llama系列模型是闭源gpt3.5火爆之后开源的强大经典decoder-only模型，后面诞生的诸多llm多多少少都带有LLama的影子。
LLama2模型结构的主要改进 1. 将Layer_norm更换为RMS_norm 在NLP模型中，归一化对模型训练过程中防止loss起飞有重要作用。
经典layer norm计算公式如下： $$ y = \frac {x - E(x)} {\sqrt{Var(x) &#43; \epsilon}} * \gamma &#43; \beta $$ $$ E(x) = \frac {1} {N} \sum^{N}_{i=0} {x_i} $$
$$ Var(x) = \frac {1} {N} \sum^{N}_{i=0}{(x_i - E(x))^2} $$
其中,$γ$和$β$是可学习的参数。分母加上一个极小的数$ε$防止分母为0。
RMS norm其实是layer norm的变体，为了加快计算，省去了求均值的过程，也删除了偏置值$β$。 $$ y = \frac {x} { \sqrt {Mean(x^2) &#43; \epsilon}} * \gamma $$
$$ Mean(x^2) = \frac {1} {N} \sum^{N}_{i=0}({x_i}^2) $$ $γ$是可学习的参数
1 2 3 4 5 6 7 8 9 10 11 12 13 class RMSNorm(torch.'><title>LLama2模型结构方面的改进</title>

<link rel='canonical' href='https://charent.github.io/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/'>

<link rel="stylesheet" href="/scss/style.min.65463afd28d606277b44441c8bbb8b0277823a2d0c03ab8ba9d0567d274f7b43.css"><meta property='og:title' content='LLama2模型结构方面的改进'>
<meta property='og:description' content='LLama系列模型 llama系列模型是闭源gpt3.5火爆之后开源的强大经典decoder-only模型，后面诞生的诸多llm多多少少都带有LLama的影子。
LLama2模型结构的主要改进 1. 将Layer_norm更换为RMS_norm 在NLP模型中，归一化对模型训练过程中防止loss起飞有重要作用。
经典layer norm计算公式如下： $$ y = \frac {x - E(x)} {\sqrt{Var(x) &#43; \epsilon}} * \gamma &#43; \beta $$ $$ E(x) = \frac {1} {N} \sum^{N}_{i=0} {x_i} $$
$$ Var(x) = \frac {1} {N} \sum^{N}_{i=0}{(x_i - E(x))^2} $$
其中,$γ$和$β$是可学习的参数。分母加上一个极小的数$ε$防止分母为0。
RMS norm其实是layer norm的变体，为了加快计算，省去了求均值的过程，也删除了偏置值$β$。 $$ y = \frac {x} { \sqrt {Mean(x^2) &#43; \epsilon}} * \gamma $$
$$ Mean(x^2) = \frac {1} {N} \sum^{N}_{i=0}({x_i}^2) $$ $γ$是可学习的参数
1 2 3 4 5 6 7 8 9 10 11 12 13 class RMSNorm(torch.'>
<meta property='og:url' content='https://charent.github.io/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/'>
<meta property='og:site_name' content='Charent的博客'>
<meta property='og:type' content='article'><meta property='article:section' content='Post' /><meta property='article:published_time' content='2023-11-12T00:00:00&#43;00:00'/><meta property='article:modified_time' content='2023-11-12T00:00:00&#43;00:00'/>
<meta name="twitter:title" content="LLama2模型结构方面的改进">
<meta name="twitter:description" content="LLama系列模型 llama系列模型是闭源gpt3.5火爆之后开源的强大经典decoder-only模型，后面诞生的诸多llm多多少少都带有LLama的影子。
LLama2模型结构的主要改进 1. 将Layer_norm更换为RMS_norm 在NLP模型中，归一化对模型训练过程中防止loss起飞有重要作用。
经典layer norm计算公式如下： $$ y = \frac {x - E(x)} {\sqrt{Var(x) &#43; \epsilon}} * \gamma &#43; \beta $$ $$ E(x) = \frac {1} {N} \sum^{N}_{i=0} {x_i} $$
$$ Var(x) = \frac {1} {N} \sum^{N}_{i=0}{(x_i - E(x))^2} $$
其中,$γ$和$β$是可学习的参数。分母加上一个极小的数$ε$防止分母为0。
RMS norm其实是layer norm的变体，为了加快计算，省去了求均值的过程，也删除了偏置值$β$。 $$ y = \frac {x} { \sqrt {Mean(x^2) &#43; \epsilon}} * \gamma $$
$$ Mean(x^2) = \frac {1} {N} \sum^{N}_{i=0}({x_i}^2) $$ $γ$是可学习的参数
1 2 3 4 5 6 7 8 9 10 11 12 13 class RMSNorm(torch.">
    </head>
    <body class="
    article-page
    ">
    <script>
        (function() {
            const colorSchemeKey = 'StackColorScheme';
            if(!localStorage.getItem(colorSchemeKey)){
                localStorage.setItem(colorSchemeKey, "auto");
            }
        })();
    </script><script>
    (function() {
        const colorSchemeKey = 'StackColorScheme';
        const colorSchemeItem = localStorage.getItem(colorSchemeKey);
        const supportDarkMode = window.matchMedia('(prefers-color-scheme: dark)').matches === true;

        if (colorSchemeItem == 'dark' || colorSchemeItem === 'auto' && supportDarkMode) {
            

            document.documentElement.dataset.scheme = 'dark';
        } else {
            document.documentElement.dataset.scheme = 'light';
        }
    })();
</script>
<div class="container main-container flex on-phone--column extended"><aside class="sidebar left-sidebar sticky ">
    <button class="hamburger hamburger--spin" type="button" id="toggle-menu" aria-label="切换菜单">
        <span class="hamburger-box">
            <span class="hamburger-inner"></span>
        </span>
    </button>

    <header>
        
            
            <figure class="site-avatar">
                <a href="/">
                
                    
                    
                    
                        
                        <img src="/img/avatar_hu544a9780599d84daff7ca906f17dd84c_13353_300x0_resize_box_3.png" width="300"
                            height="300" class="site-logo" loading="lazy" alt="Avatar">
                    
                
                </a>
                
                    <span class="emoji">😆</span>
                
            </figure>
            
        
        
        <div class="site-meta">
            <h1 class="site-name"><a href="/">Charent的博客</a></h1>
            <h2 class="site-description">不积硅步，无以至千里</h2>
        </div>
    </header><ol class="social-menu">
            
                <li>
                    <a 
                        href='https://github.com/charent'
                        target="_blank"
                        title="GitHub"
                    >
                        
                        
                            <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-brand-github" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
  <path d="M9 19c-4.3 1.4 -4.3 -2.5 -6 -3m12 5v-3.5c0 -1 .1 -1.4 -.5 -2c2.8 -.3 5.5 -1.4 5.5 -6a4.6 4.6 0 0 0 -1.3 -3.2a4.2 4.2 0 0 0 -.1 -3.2s-1.1 -.3 -3.5 1.3a12.3 12.3 0 0 0 -6.2 0c-2.4 -1.6 -3.5 -1.3 -3.5 -1.3a4.2 4.2 0 0 0 -.1 3.2a4.6 4.6 0 0 0 -1.3 3.2c0 4.6 2.7 5.7 5.5 6c-.6 .6 -.6 1.2 -.5 2v3.5" />
</svg>



                        
                    </a>
                </li>
            
        </ol><ol class="menu" id="main-menu">
        
        
        

        <li >
            <a href='/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-home" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <polyline points="5 12 3 12 12 3 21 12 19 12" />
  <path d="M5 12v7a2 2 0 0 0 2 2h10a2 2 0 0 0 2 -2v-7" />
  <path d="M9 21v-6a2 2 0 0 1 2 -2h2a2 2 0 0 1 2 2v6" />
</svg>



                
                <span>主页</span>
            </a>
        </li>
        
        

        <li >
            <a href='/archives/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-archive" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <rect x="3" y="4" width="18" height="4" rx="2" />
  <path d="M5 8v10a2 2 0 0 0 2 2h10a2 2 0 0 0 2 -2v-10" />
  <line x1="10" y1="12" x2="14" y2="12" />
</svg>



                
                <span>归档</span>
            </a>
        </li>
        
        

        <li >
            <a href='/search/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-search" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <circle cx="10" cy="10" r="7" />
  <line x1="21" y1="21" x2="15" y2="15" />
</svg>



                
                <span>搜索</span>
            </a>
        </li>
        
        

        <li >
            <a href='/%E5%85%B3%E4%BA%8E/' >
                
                
                
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-link" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <path d="M10 14a3.5 3.5 0 0 0 5 0l4 -4a3.5 3.5 0 0 0 -5 -5l-.5 .5" />
  <path d="M14 10a3.5 3.5 0 0 0 -5 0l-4 4a3.5 3.5 0 0 0 5 5l.5 -.5" />
</svg>



                
                <span>关于</span>
            </a>
        </li>
        

        <div class="menu-bottom-section">
                <li id="i18n-switch">  
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-language" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z" fill="none"/>
  <path d="M4 5h7" />
  <path d="M9 3v2c0 4.418 -2.239 8 -5 8" />
  <path d="M5 9c-.003 2.144 2.952 3.908 6.7 4" />
  <path d="M12 20l4 -9l4 9" />
  <path d="M19.1 18h-6.2" />
</svg>



                    <select name="language" onchange="window.location.href = this.selectedOptions[0].value">
                        
                            <option value="https://charent.github.io/" selected>中文</option>
                        
                            <option value="https://charent.github.io/en/" >English</option>
                        
                    </select>
                </li>
            
            
            
                <li id="dark-mode-toggle">
                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-toggle-left" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <circle cx="8" cy="12" r="2" />
  <rect x="2" y="6" width="20" height="12" rx="6" />
</svg>



                    <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-toggle-right" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <circle cx="16" cy="12" r="2" />
  <rect x="2" y="6" width="20" height="12" rx="6" />
</svg>



                    <span>暗色模式</span>
                </li>
            
        </div>
    </ol>
</aside>
<main class="main full-width">
    <article class="main-article">
    <header class="article-header">

    <div class="article-details">
    
    <header class="article-category">
        
            <a href="/categories/llm%E5%A4%A7%E6%A8%A1%E5%9E%8B/" >
                llm大模型
            </a>
        
    </header>
    

    <div class="article-title-wrapper">
        <h2 class="article-title">
            <a href="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/">LLama2模型结构方面的改进</a>
        </h2>
    
        
    </div>

    
    <footer class="article-time">
        
            <div>
                <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-calendar-time" width="56" height="56" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <path d="M11.795 21h-6.795a2 2 0 0 1 -2 -2v-12a2 2 0 0 1 2 -2h12a2 2 0 0 1 2 2v4" />
  <circle cx="18" cy="18" r="4" />
  <path d="M15 3v4" />
  <path d="M7 3v4" />
  <path d="M3 11h16" />
  <path d="M18 16.496v1.504l1 1" />
</svg>
                <time class="article-time--published">2023-11-12</time>
            </div>
        

        
    </footer>
    

    
</div>
</header>

    <section class="article-content">
    
    
    <h1 id="llama系列模型">LLama系列模型</h1>
<p>llama系列模型是闭源gpt3.5火爆之后开源的强大经典decoder-only模型，后面诞生的诸多llm多多少少都带有LLama的影子。</p>
<h1 id="llama2模型结构的主要改进">LLama2模型结构的主要改进</h1>
<h2 id="1-将layer_norm更换为rms_norm">1. 将Layer_norm更换为RMS_norm</h2>
<p>在NLP模型中，归一化对模型训练过程中防止loss起飞有重要作用。</p>
<ul>
<li>经典layer norm计算公式如下：
$$
y = \frac {x - E(x)} {\sqrt{Var(x) + \epsilon}} * \gamma + \beta
$$</li>
</ul>
<p>$$
E(x) = \frac {1} {N} \sum^{N}_{i=0} {x_i}
$$</p>
<p>$$
Var(x) = \frac {1} {N} \sum^{N}_{i=0}{(x_i - E(x))^2}
$$</p>
<p>其中,$γ$和$β$是可学习的参数。分母加上一个极小的数$ε$防止分母为0。</p>
<ul>
<li>RMS norm其实是layer norm的变体，为了加快计算，省去了求均值的过程，也删除了偏置值$β$。</li>
</ul>
<p>$$
y = \frac {x} { \sqrt {Mean(x^2) + \epsilon}} * \gamma
$$</p>
<p>$$
Mean(x^2) = \frac {1} {N} \sum^{N}_{i=0}({x_i}^2)
$$
$γ$是可学习的参数</p>
<div class="highlight"><div class="chroma">
<table class="lntable"><tr><td class="lntd">
<pre tabindex="0" class="chroma"><code><span class="lnt"> 1
</span><span class="lnt"> 2
</span><span class="lnt"> 3
</span><span class="lnt"> 4
</span><span class="lnt"> 5
</span><span class="lnt"> 6
</span><span class="lnt"> 7
</span><span class="lnt"> 8
</span><span class="lnt"> 9
</span><span class="lnt">10
</span><span class="lnt">11
</span><span class="lnt">12
</span><span class="lnt">13
</span></code></pre></td>
<td class="lntd">
<pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">RMSNorm</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-6</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">eps</span> <span class="o">=</span> <span class="n">eps</span>  <span class="c1"># ε</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">gama</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">dim</span><span class="p">))</span> <span class="c1">#可学习参数γ</span>
</span></span><span class="line"><span class="cl"><span class="err">​</span>
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="nf">_norm</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="c1"># RMSNorm</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="n">x</span> <span class="o">*</span> <span class="n">torch</span><span class="o">.</span><span class="n">rsqrt</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">eps</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="err">​</span>
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="n">output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_norm</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">float</span><span class="p">())</span><span class="o">.</span><span class="n">type_as</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="n">output</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">gama</span>
</span></span></code></pre></td></tr></table>
</div>
</div><h2 id="2-q在与k相乘之前先使用rope进行位置编码">2. Q在与K相乘之前，先使用RoPE进行位置编码</h2>
<p><img src="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/rope.png"
	width="1072"
	height="556"
	srcset="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/rope_hua244ea9e08212bfb4ed03f83d4a6738d_87886_480x0_resize_box_3.png 480w, /p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/rope_hua244ea9e08212bfb4ed03f83d4a6738d_87886_1024x0_resize_box_3.png 1024w"
	loading="lazy"
	
		alt="rope"
	
	
		class="gallery-image" 
		data-flex-grow="192"
		data-flex-basis="462px"
	
></p>
<div class="highlight"><div class="chroma">
<table class="lntable"><tr><td class="lntd">
<pre tabindex="0" class="chroma"><code><span class="lnt"> 1
</span><span class="lnt"> 2
</span><span class="lnt"> 3
</span><span class="lnt"> 4
</span><span class="lnt"> 5
</span><span class="lnt"> 6
</span><span class="lnt"> 7
</span><span class="lnt"> 8
</span><span class="lnt"> 9
</span><span class="lnt">10
</span><span class="lnt">11
</span><span class="lnt">12
</span><span class="lnt">13
</span><span class="lnt">14
</span><span class="lnt">15
</span><span class="lnt">16
</span><span class="lnt">17
</span><span class="lnt">18
</span><span class="lnt">19
</span><span class="lnt">20
</span><span class="lnt">21
</span><span class="lnt">22
</span><span class="lnt">23
</span><span class="lnt">24
</span><span class="lnt">25
</span><span class="lnt">26
</span><span class="lnt">27
</span><span class="lnt">28
</span><span class="lnt">29
</span><span class="lnt">30
</span><span class="lnt">31
</span><span class="lnt">32
</span><span class="lnt">33
</span><span class="lnt">34
</span><span class="lnt">35
</span><span class="lnt">36
</span><span class="lnt">37
</span><span class="lnt">38
</span><span class="lnt">39
</span><span class="lnt">40
</span><span class="lnt">41
</span><span class="lnt">42
</span><span class="lnt">43
</span><span class="lnt">44
</span><span class="lnt">45
</span><span class="lnt">46
</span><span class="lnt">47
</span><span class="lnt">48
</span><span class="lnt">49
</span><span class="lnt">50
</span><span class="lnt">51
</span><span class="lnt">52
</span><span class="lnt">53
</span><span class="lnt">54
</span><span class="lnt">55
</span><span class="lnt">56
</span><span class="lnt">57
</span><span class="lnt">58
</span><span class="lnt">59
</span><span class="lnt">60
</span><span class="lnt">61
</span><span class="lnt">62
</span><span class="lnt">63
</span><span class="lnt">64
</span><span class="lnt">65
</span><span class="lnt">66
</span><span class="lnt">67
</span><span class="lnt">68
</span><span class="lnt">69
</span><span class="lnt">70
</span><span class="lnt">71
</span><span class="lnt">72
</span><span class="lnt">73
</span></code></pre></td>
<td class="lntd">
<pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="c1"># 作者：CodeLearner</span>
</span></span><span class="line"><span class="cl"><span class="c1"># 链接：https://zhuanlan.zhihu.com/p/649756898</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="k">def</span> <span class="nf">precompute_freqs_cis</span><span class="p">(</span><span class="n">dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">end</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">theta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">10000.0</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># 计算词向量元素两两分组以后，每组元素对应的旋转角度 </span>
</span></span><span class="line"><span class="cl">    <span class="c1"># arange生成[0,2,4...126]</span>
</span></span><span class="line"><span class="cl">    <span class="n">freqs</span> <span class="o">=</span> <span class="mf">1.0</span> <span class="o">/</span> <span class="p">(</span><span class="n">theta</span> <span class="o">**</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">dim</span><span class="p">,</span> <span class="mi">2</span><span class="p">)[:</span> <span class="p">(</span><span class="n">dim</span> <span class="o">//</span> <span class="mi">2</span><span class="p">)]</span><span class="o">.</span><span class="n">float</span><span class="p">()</span> <span class="o">/</span> <span class="n">dim</span><span class="p">))</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># t = [0,....end]</span>
</span></span><span class="line"><span class="cl">    <span class="n">t</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">end</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">freqs</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>  <span class="c1"># type: ignore</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># t为列向量 freqs为行向量做外积</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># freqs.shape = (t.len(),freqs.len()) #shape (end,dim//2)</span>
</span></span><span class="line"><span class="cl">    <span class="n">freqs</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">outer</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">freqs</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span>  <span class="c1"># type: ignore</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># 生成复数</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># torch.polar(abs,angle) -&gt; abs*cos(angle) + abs*sin(angle)*j</span>
</span></span><span class="line"><span class="cl">    <span class="n">freqs_cis</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">polar</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">freqs</span><span class="p">),</span> <span class="n">freqs</span><span class="p">)</span>  <span class="c1"># complex64</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># freqs_cis.shape  = (end,dim//2)</span>
</span></span><span class="line"><span class="cl">    <span class="k">return</span> <span class="n">freqs_cis</span>
</span></span><span class="line"><span class="cl"><span class="err">​</span>
</span></span><span class="line"><span class="cl"><span class="k">def</span> <span class="nf">reshape_for_broadcast</span><span class="p">(</span><span class="n">freqs_cis</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># ndim为x的维度数 ,此时应该为4</span>
</span></span><span class="line"><span class="cl">    <span class="n">ndim</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">ndim</span>
</span></span><span class="line"><span class="cl">    <span class="k">assert</span> <span class="mi">0</span> <span class="o">&lt;=</span> <span class="mi">1</span> <span class="o">&lt;</span> <span class="n">ndim</span>
</span></span><span class="line"><span class="cl">    <span class="k">assert</span> <span class="n">freqs_cis</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
</span></span><span class="line"><span class="cl">    <span class="n">shape</span> <span class="o">=</span> <span class="p">[</span><span class="n">d</span> <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">i</span> <span class="o">==</span> <span class="n">ndim</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">else</span> <span class="mi">1</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">d</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">shape</span><span class="p">)]</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># (1,x.shape[1],1,x.shape[-1])</span>
</span></span><span class="line"><span class="cl">    <span class="k">return</span> <span class="n">freqs_cis</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">*</span><span class="n">shape</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="err">​</span>
</span></span><span class="line"><span class="cl"><span class="k">def</span> <span class="nf">apply_rotary_emb</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">    <span class="n">xq</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">    <span class="n">xk</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">    <span class="n">freqs_cis</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
</span></span><span class="line"><span class="cl"><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># xq.shape = [bsz, seqlen, self.n_local_heads, self.head_dim]</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># xq_.shape = [bsz, seqlen, self.n_local_heads, self.head_dim//2 , 2]</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># torch.view_as_complex用于将二维向量转换为复数域 torch.view_as_complex即([x,y]) -&gt; (x+yj)</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># 所以经过view_as_complex变换后xq_.shape = [bsz, seqlen, self.n_local_heads, self.head_dim//2]</span>
</span></span><span class="line"><span class="cl">    <span class="n">xq_</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_complex</span><span class="p">(</span><span class="n">xq</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">*</span><span class="n">xq</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
</span></span><span class="line"><span class="cl">    <span class="n">xk_</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_complex</span><span class="p">(</span><span class="n">xk</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">*</span><span class="n">xk</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
</span></span><span class="line"><span class="cl">    
</span></span><span class="line"><span class="cl">    
</span></span><span class="line"><span class="cl">    <span class="n">freqs_cis</span> <span class="o">=</span> <span class="n">reshape_for_broadcast</span><span class="p">(</span><span class="n">freqs_cis</span><span class="p">,</span> <span class="n">xq_</span><span class="p">)</span> <span class="c1"># freqs_cis.shape = (1,x.shape[1],1,x.shape[-1])</span>
</span></span><span class="line"><span class="cl">    
</span></span><span class="line"><span class="cl">    <span class="c1"># xq_ 与freqs_cis广播哈达玛积</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># [bsz, seqlen, self.n_local_heads, self.head_dim//2] * [1,seqlen,1,self.head_dim//2]</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># torch.view_as_real用于将复数再转换回实数向量, 再经过flatten展平第4个维度 </span>
</span></span><span class="line"><span class="cl">    <span class="c1"># [bsz, seqlen, self.n_local_heads, self.head_dim//2] -&gt;[bsz, seqlen, self.n_local_heads, self.head_dim//2,2 ] -&gt;[bsz, seqlen, self.n_local_heads, self.head_dim]</span>
</span></span><span class="line"><span class="cl">    <span class="n">xq_out</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">xq_</span> <span class="o">*</span> <span class="n">freqs_cis</span><span class="p">)</span><span class="o">.</span><span class="n">flatten</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">    <span class="n">xk_out</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">xk_</span> <span class="o">*</span> <span class="n">freqs_cis</span><span class="p">)</span><span class="o">.</span><span class="n">flatten</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">    <span class="k">return</span> <span class="n">xq_out</span><span class="o">.</span><span class="n">type_as</span><span class="p">(</span><span class="n">xq</span><span class="p">),</span> <span class="n">xk_out</span><span class="o">.</span><span class="n">type_as</span><span class="p">(</span><span class="n">xk</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="c1"># 精简版Attention</span>
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">Attention</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">args</span><span class="p">:</span> <span class="n">ModelArgs</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wq</span> <span class="o">=</span> <span class="n">Linear</span><span class="p">(</span><span class="o">...</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wk</span> <span class="o">=</span> <span class="n">Linear</span><span class="p">(</span><span class="o">...</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wv</span> <span class="o">=</span> <span class="n">Linear</span><span class="p">(</span><span class="o">...</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">freqs_cis</span> <span class="o">=</span> <span class="n">precompute_freqs_cis</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">max_seq_len</span> <span class="o">*</span> <span class="mi">2</span><span class="p">)</span>
</span></span><span class="line"><span class="cl"><span class="err">​</span>
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span>
</span></span><span class="line"><span class="cl">        <span class="n">xq</span><span class="p">,</span> <span class="n">xk</span><span class="p">,</span> <span class="n">xv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">wq</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">wk</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">wv</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">xq</span> <span class="o">=</span> <span class="n">xq</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_heads</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">xk</span> <span class="o">=</span> <span class="n">xk</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">xv</span> <span class="o">=</span> <span class="n">xv</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">         <span class="c1"># attention 操作之前，应用旋转位置编码</span>
</span></span><span class="line"><span class="cl">        <span class="n">xq</span><span class="p">,</span> <span class="n">xk</span> <span class="o">=</span> <span class="n">apply_rotary_emb</span><span class="p">(</span><span class="n">xq</span><span class="p">,</span> <span class="n">xk</span><span class="p">,</span> <span class="n">freqs_cis</span><span class="o">=</span><span class="n">freqs_cis</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="c1">#...</span>
</span></span><span class="line"><span class="cl">        <span class="c1"># 进行后续Attention计算</span>
</span></span><span class="line"><span class="cl">        <span class="n">scores</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">xq</span><span class="p">,</span> <span class="n">xk</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span> <span class="o">/</span> <span class="n">math</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">scores</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">scores</span><span class="o">.</span><span class="n">float</span><span class="p">(),</span> <span class="n">dim</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">output</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">scores</span><span class="p">,</span> <span class="n">xv</span><span class="p">)</span>  <span class="c1"># (batch_size, seq_len, dim)</span>
</span></span><span class="line"><span class="cl">  <span class="c1"># ......</span>
</span></span></code></pre></td></tr></table>
</div>
</div><h2 id="3-引入kv-cache并采用group-query-attention">3. 引入KV Cache，并采用Group Query Attention</h2>
<p><img src="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/kv-cache.jpg"
	width="1566"
	height="1054"
	srcset="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/kv-cache_hu506057db48b1cfc262d6b7344f6873c4_145994_480x0_resize_q75_box.jpg 480w, /p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/kv-cache_hu506057db48b1cfc262d6b7344f6873c4_145994_1024x0_resize_q75_box.jpg 1024w"
	loading="lazy"
	
		alt="kv-cache"
	
	
		class="gallery-image" 
		data-flex-grow="148"
		data-flex-basis="356px"
	
>
出处见图片水印。</p>
<div class="highlight"><div class="chroma">
<table class="lntable"><tr><td class="lntd">
<pre tabindex="0" class="chroma"><code><span class="lnt"> 1
</span><span class="lnt"> 2
</span><span class="lnt"> 3
</span><span class="lnt"> 4
</span><span class="lnt"> 5
</span><span class="lnt"> 6
</span><span class="lnt"> 7
</span><span class="lnt"> 8
</span><span class="lnt"> 9
</span><span class="lnt">10
</span><span class="lnt">11
</span><span class="lnt">12
</span><span class="lnt">13
</span><span class="lnt">14
</span><span class="lnt">15
</span><span class="lnt">16
</span><span class="lnt">17
</span></code></pre></td>
<td class="lntd">
<pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="c1"># 作者：CodeLearner</span>
</span></span><span class="line"><span class="cl"><span class="c1"># 链接：https://zhuanlan.zhihu.com/p/649756898</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="k">def</span> <span class="nf">mha</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">c_attn</span><span class="p">,</span> <span class="n">c_proj</span><span class="p">,</span> <span class="n">n_head</span><span class="p">,</span> <span class="n">kvcache</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># [n_seq, n_embd] -&gt; [n_seq, n_embd]</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># qkv projection</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># when we pass kvcache, n_seq = 1. so we will compute new_q, new_k and new_v</span>
</span></span><span class="line"><span class="cl">    <span class="n">x</span> <span class="o">=</span> <span class="n">linear</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="o">**</span><span class="n">c_attn</span><span class="p">)</span>  <span class="c1"># [n_seq, n_embd] -&gt; [n_seq, 3*n_embd]</span>
</span></span><span class="line"><span class="cl">    <span class="c1"># split into qkv</span>
</span></span><span class="line"><span class="cl">    <span class="n">qkv</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">axis</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># [n_seq, 3*n_embd] -&gt; [3, n_seq, n_embd]</span>
</span></span><span class="line"><span class="cl">    <span class="k">if</span> <span class="n">kvcache</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">        <span class="c1"># qkv</span>
</span></span><span class="line"><span class="cl">        <span class="n">new_q</span><span class="p">,</span> <span class="n">new_k</span><span class="p">,</span> <span class="n">new_v</span> <span class="o">=</span> <span class="n">qkv</span>  <span class="c1"># new_q, new_k, new_v = [1, n_embd]</span>
</span></span><span class="line"><span class="cl">        <span class="n">old_k</span><span class="p">,</span> <span class="n">old_v</span> <span class="o">=</span> <span class="n">kvcache</span>
</span></span><span class="line"><span class="cl">        <span class="n">k</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">([</span><span class="n">old_k</span><span class="p">,</span> <span class="n">new_k</span><span class="p">])</span> <span class="c1"># k = [n_seq, n_embd], where n_seq = prev_n_seq + 1</span>
</span></span><span class="line"><span class="cl">        <span class="n">v</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">([</span><span class="n">old_v</span><span class="p">,</span> <span class="n">new_v</span><span class="p">])</span> <span class="c1"># v = [n_seq, n_embd], where n_seq = prev_n_seq + 1</span>
</span></span><span class="line"><span class="cl">        <span class="n">qkv</span> <span class="o">=</span> <span class="p">[</span><span class="n">new_q</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>
</span></span></code></pre></td></tr></table>
</div>
</div><h2 id="4-用silu激活函数代替relugelu">4. 用SiLU激活函数代替RELU/GELU</h2>
<p><img src="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/act_fun.png"
	width="720"
	height="561"
	srcset="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/act_fun_hu1bf631b7966cd0dd3d56cd597aac1f28_72453_480x0_resize_box_3.png 480w, /p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/act_fun_hu1bf631b7966cd0dd3d56cd597aac1f28_72453_1024x0_resize_box_3.png 1024w"
	loading="lazy"
	
		alt="act fun"
	
	
		class="gallery-image" 
		data-flex-grow="128"
		data-flex-basis="308px"
	
>
这里先介绍Swish激活函数，计算公式如下：
$$
Swish(x) = x · \sigma (\beta x)
$$</p>
<p>$$
\sigma (x) = \frac{1} {1 + e^{-x} }
$$</p>
<p>其中，$σ$是<code>Sigmoid</code>函数， $\beta$ 是一个参数，上图为$β=0.5$时的函数图像。当$β=1$时，Swish函数就是SiLU函数。</p>
<h2 id="5-使用分组查询注意力-grouped-query-attention-gqa">5. 使用分组查询注意力 (Grouped Query Attention, GQA)</h2>
<p><img src="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/maq_gqa.png"
	width="1080"
	height="349"
	srcset="/p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/maq_gqa_hu4b2c50f6cc9ed5691ca0aeb990bfde87_51507_480x0_resize_box_3.png 480w, /p/llama2%E6%A8%A1%E5%9E%8B%E7%BB%93%E6%9E%84%E6%96%B9%E9%9D%A2%E7%9A%84%E6%94%B9%E8%BF%9B/maq_gqa_hu4b2c50f6cc9ed5691ca0aeb990bfde87_51507_1024x0_resize_box_3.png 1024w"
	loading="lazy"
	
		alt="maq_gqa"
	
	
		class="gallery-image" 
		data-flex-grow="309"
		data-flex-basis="742px"
	
></p>
<div class="highlight"><div class="chroma">
<table class="lntable"><tr><td class="lntd">
<pre tabindex="0" class="chroma"><code><span class="lnt">  1
</span><span class="lnt">  2
</span><span class="lnt">  3
</span><span class="lnt">  4
</span><span class="lnt">  5
</span><span class="lnt">  6
</span><span class="lnt">  7
</span><span class="lnt">  8
</span><span class="lnt">  9
</span><span class="lnt"> 10
</span><span class="lnt"> 11
</span><span class="lnt"> 12
</span><span class="lnt"> 13
</span><span class="lnt"> 14
</span><span class="lnt"> 15
</span><span class="lnt"> 16
</span><span class="lnt"> 17
</span><span class="lnt"> 18
</span><span class="lnt"> 19
</span><span class="lnt"> 20
</span><span class="lnt"> 21
</span><span class="lnt"> 22
</span><span class="lnt"> 23
</span><span class="lnt"> 24
</span><span class="lnt"> 25
</span><span class="lnt"> 26
</span><span class="lnt"> 27
</span><span class="lnt"> 28
</span><span class="lnt"> 29
</span><span class="lnt"> 30
</span><span class="lnt"> 31
</span><span class="lnt"> 32
</span><span class="lnt"> 33
</span><span class="lnt"> 34
</span><span class="lnt"> 35
</span><span class="lnt"> 36
</span><span class="lnt"> 37
</span><span class="lnt"> 38
</span><span class="lnt"> 39
</span><span class="lnt"> 40
</span><span class="lnt"> 41
</span><span class="lnt"> 42
</span><span class="lnt"> 43
</span><span class="lnt"> 44
</span><span class="lnt"> 45
</span><span class="lnt"> 46
</span><span class="lnt"> 47
</span><span class="lnt"> 48
</span><span class="lnt"> 49
</span><span class="lnt"> 50
</span><span class="lnt"> 51
</span><span class="lnt"> 52
</span><span class="lnt"> 53
</span><span class="lnt"> 54
</span><span class="lnt"> 55
</span><span class="lnt"> 56
</span><span class="lnt"> 57
</span><span class="lnt"> 58
</span><span class="lnt"> 59
</span><span class="lnt"> 60
</span><span class="lnt"> 61
</span><span class="lnt"> 62
</span><span class="lnt"> 63
</span><span class="lnt"> 64
</span><span class="lnt"> 65
</span><span class="lnt"> 66
</span><span class="lnt"> 67
</span><span class="lnt"> 68
</span><span class="lnt"> 69
</span><span class="lnt"> 70
</span><span class="lnt"> 71
</span><span class="lnt"> 72
</span><span class="lnt"> 73
</span><span class="lnt"> 74
</span><span class="lnt"> 75
</span><span class="lnt"> 76
</span><span class="lnt"> 77
</span><span class="lnt"> 78
</span><span class="lnt"> 79
</span><span class="lnt"> 80
</span><span class="lnt"> 81
</span><span class="lnt"> 82
</span><span class="lnt"> 83
</span><span class="lnt"> 84
</span><span class="lnt"> 85
</span><span class="lnt"> 86
</span><span class="lnt"> 87
</span><span class="lnt"> 88
</span><span class="lnt"> 89
</span><span class="lnt"> 90
</span><span class="lnt"> 91
</span><span class="lnt"> 92
</span><span class="lnt"> 93
</span><span class="lnt"> 94
</span><span class="lnt"> 95
</span><span class="lnt"> 96
</span><span class="lnt"> 97
</span><span class="lnt"> 98
</span><span class="lnt"> 99
</span><span class="lnt">100
</span><span class="lnt">101
</span><span class="lnt">102
</span><span class="lnt">103
</span><span class="lnt">104
</span><span class="lnt">105
</span><span class="lnt">106
</span><span class="lnt">107
</span><span class="lnt">108
</span><span class="lnt">109
</span></code></pre></td>
<td class="lntd">
<pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="c1"># copied from llama2</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="k">def</span> <span class="nf">repeat_kv</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">n_rep</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">    <span class="s2">&#34;&#34;&#34;torch.repeat_interleave(x, dim=2, repeats=n_rep)&#34;&#34;&#34;</span>
</span></span><span class="line"><span class="cl">    <span class="n">bs</span><span class="p">,</span> <span class="n">slen</span><span class="p">,</span> <span class="n">n_kv_heads</span><span class="p">,</span> <span class="n">head_dim</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span>
</span></span><span class="line"><span class="cl">    <span class="k">if</span> <span class="n">n_rep</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span> <span class="c1"># MHA</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="n">x</span>
</span></span><span class="line"><span class="cl">    <span class="k">return</span> <span class="p">(</span> <span class="c1"># MQA / GQA</span>
</span></span><span class="line"><span class="cl">        <span class="n">x</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">:,</span> <span class="kc">None</span><span class="p">,</span> <span class="p">:]</span>
</span></span><span class="line"><span class="cl">        <span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">slen</span><span class="p">,</span> <span class="n">n_kv_heads</span><span class="p">,</span> <span class="n">n_rep</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">slen</span><span class="p">,</span> <span class="n">n_kv_heads</span> <span class="o">*</span> <span class="n">n_rep</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">    <span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl"><span class="k">class</span> <span class="nc">Attention</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">args</span><span class="p">:</span> <span class="n">ModelArgs</span><span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">n_kv_heads</span> <span class="o">=</span> <span class="n">args</span><span class="o">.</span><span class="n">n_heads</span> <span class="k">if</span> <span class="n">args</span><span class="o">.</span><span class="n">n_kv_heads</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">args</span><span class="o">.</span><span class="n">n_kv_heads</span>
</span></span><span class="line"><span class="cl">        <span class="n">model_parallel_size</span> <span class="o">=</span> <span class="n">fs_init</span><span class="o">.</span><span class="n">get_model_parallel_world_size</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">n_local_heads</span> <span class="o">=</span> <span class="n">args</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">//</span> <span class="n">model_parallel_size</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_kv_heads</span> <span class="o">//</span> <span class="n">model_parallel_size</span> <span class="c1"># 此处 </span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">n_rep</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_heads</span> <span class="o">//</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span> <span class="c1"># 此处 几个组</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span> <span class="o">=</span> <span class="n">args</span><span class="o">.</span><span class="n">dim</span> <span class="o">//</span> <span class="n">args</span><span class="o">.</span><span class="n">n_heads</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wq</span> <span class="o">=</span> <span class="n">ColumnParallelLinear</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">            <span class="n">args</span><span class="o">.</span><span class="n">dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">args</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">gather_output</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">init_method</span><span class="o">=</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wk</span> <span class="o">=</span> <span class="n">ColumnParallelLinear</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">            <span class="n">args</span><span class="o">.</span><span class="n">dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="bp">self</span><span class="o">.</span><span class="n">n_kv_heads</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span> <span class="c1"># 初始化为单个组内的一份</span>
</span></span><span class="line"><span class="cl">            <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">gather_output</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">init_method</span><span class="o">=</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wv</span> <span class="o">=</span> <span class="n">ColumnParallelLinear</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">            <span class="n">args</span><span class="o">.</span><span class="n">dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="bp">self</span><span class="o">.</span><span class="n">n_kv_heads</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span> <span class="c1"># # 初始化为单个组内的一份</span>
</span></span><span class="line"><span class="cl">            <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">gather_output</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">init_method</span><span class="o">=</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">wo</span> <span class="o">=</span> <span class="n">RowParallelLinear</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">            <span class="n">args</span><span class="o">.</span><span class="n">n_heads</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">args</span><span class="o">.</span><span class="n">dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">input_is_parallel</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="n">init_method</span><span class="o">=</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">cache_k</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">            <span class="p">(</span>
</span></span><span class="line"><span class="cl">                <span class="n">args</span><span class="o">.</span><span class="n">max_batch_size</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="n">args</span><span class="o">.</span><span class="n">max_seq_len</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="p">)</span><span class="o">.</span><span class="n">cuda</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">cache_v</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">            <span class="p">(</span>
</span></span><span class="line"><span class="cl">                <span class="n">args</span><span class="o">.</span><span class="n">max_batch_size</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="n">args</span><span class="o">.</span><span class="n">max_seq_len</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">                <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">            <span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="p">)</span><span class="o">.</span><span class="n">cuda</span><span class="p">()</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="n">x</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="n">start_pos</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="n">freqs_cis</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
</span></span><span class="line"><span class="cl">        <span class="n">mask</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span>
</span></span><span class="line"><span class="cl">    <span class="p">):</span>
</span></span><span class="line"><span class="cl">        <span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span>
</span></span><span class="line"><span class="cl">        <span class="n">xq</span><span class="p">,</span> <span class="n">xk</span><span class="p">,</span> <span class="n">xv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">wq</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">wk</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">wv</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">xq</span> <span class="o">=</span> <span class="n">xq</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_heads</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">xk</span> <span class="o">=</span> <span class="n">xk</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">xv</span> <span class="o">=</span> <span class="n">xv</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_local_kv_heads</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">xq</span><span class="p">,</span> <span class="n">xk</span> <span class="o">=</span> <span class="n">apply_rotary_emb</span><span class="p">(</span><span class="n">xq</span><span class="p">,</span> <span class="n">xk</span><span class="p">,</span> <span class="n">freqs_cis</span><span class="o">=</span><span class="n">freqs_cis</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">cache_k</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache_k</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xq</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">cache_v</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache_v</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xq</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">cache_k</span><span class="p">[:</span><span class="n">bsz</span><span class="p">,</span> <span class="n">start_pos</span> <span class="p">:</span> <span class="n">start_pos</span> <span class="o">+</span> <span class="n">seqlen</span><span class="p">]</span> <span class="o">=</span> <span class="n">xk</span>
</span></span><span class="line"><span class="cl">        <span class="bp">self</span><span class="o">.</span><span class="n">cache_v</span><span class="p">[:</span><span class="n">bsz</span><span class="p">,</span> <span class="n">start_pos</span> <span class="p">:</span> <span class="n">start_pos</span> <span class="o">+</span> <span class="n">seqlen</span><span class="p">]</span> <span class="o">=</span> <span class="n">xv</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">keys</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache_k</span><span class="p">[:</span><span class="n">bsz</span><span class="p">,</span> <span class="p">:</span> <span class="n">start_pos</span> <span class="o">+</span> <span class="n">seqlen</span><span class="p">]</span>
</span></span><span class="line"><span class="cl">        <span class="n">values</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache_v</span><span class="p">[:</span><span class="n">bsz</span><span class="p">,</span> <span class="p">:</span> <span class="n">start_pos</span> <span class="o">+</span> <span class="n">seqlen</span><span class="p">]</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="c1"># repeat k/v heads if n_kv_heads &lt; n_heads # 单个组扩展为完整head</span>
</span></span><span class="line"><span class="cl">        <span class="n">keys</span> <span class="o">=</span> <span class="n">repeat_kv</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_rep</span><span class="p">)</span>  <span class="c1"># (bs, seqlen, n_local_heads, head_dim)</span>
</span></span><span class="line"><span class="cl">        <span class="n">values</span> <span class="o">=</span> <span class="n">repeat_kv</span><span class="p">(</span><span class="n">values</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_rep</span><span class="p">)</span>  <span class="c1"># (bs, seqlen, n_local_heads, head_dim)</span>
</span></span><span class="line"><span class="cl">
</span></span><span class="line"><span class="cl">        <span class="n">xq</span> <span class="o">=</span> <span class="n">xq</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>  <span class="c1"># (bs, n_local_heads, seqlen, head_dim)</span>
</span></span><span class="line"><span class="cl">        <span class="n">keys</span> <span class="o">=</span> <span class="n">keys</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">values</span> <span class="o">=</span> <span class="n">values</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">scores</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">xq</span><span class="p">,</span> <span class="n">keys</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">))</span> <span class="o">/</span> <span class="n">math</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">head_dim</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="k">if</span> <span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
</span></span><span class="line"><span class="cl">            <span class="n">scores</span> <span class="o">=</span> <span class="n">scores</span> <span class="o">+</span> <span class="n">mask</span>  <span class="c1"># (bs, n_local_heads, seqlen, cache_len + seqlen)</span>
</span></span><span class="line"><span class="cl">        <span class="n">scores</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">scores</span><span class="o">.</span><span class="n">float</span><span class="p">(),</span> <span class="n">dim</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">type_as</span><span class="p">(</span><span class="n">xq</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="n">output</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">scores</span><span class="p">,</span> <span class="n">values</span><span class="p">)</span>  <span class="c1"># (bs, n_local_heads, seqlen, head_dim)</span>
</span></span><span class="line"><span class="cl">        <span class="n">output</span> <span class="o">=</span> <span class="n">output</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">seqlen</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
</span></span><span class="line"><span class="cl">        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">wo</span><span class="p">(</span><span class="n">output</span><span class="p">)</span>
</span></span></code></pre></td></tr></table>
</div>
</div>
</section>


    <footer class="article-footer">
    

    </footer>


    
        <link 
                rel="stylesheet" 
                href="https://cdn.jsdelivr.net/npm/katex@0.13.13/dist/katex.min.css"integrity="sha384-RZU/ijkSsFbcmivfdRBQDtwuwVqK7GMOw6IMvKyeWL2K5UAlyp6WonmB8m7Jd0Hn"crossorigin="anonymous"
            ><script 
                src="https://cdn.jsdelivr.net/npm/katex@0.13.13/dist/katex.min.js"integrity="sha384-pK1WpvzWVBQiP0/GjnvRxV4mOb0oxFuyRxJlk6vVw146n3egcN5C925NCP7a7BY8"crossorigin="anonymous"
                defer
                >
            </script><script 
                src="https://cdn.jsdelivr.net/npm/katex@0.13.13/dist/contrib/auto-render.min.js"integrity="sha384-vZTG03m&#43;2yp6N6BNi5iM4rW4oIwk5DfcNdFfxkk9ZWpDriOkXX8voJBFrAO7MpVl"crossorigin="anonymous"
                defer
                >
            </script><script>
    window.addEventListener("DOMContentLoaded", () => {
        renderMathInElement(document.querySelector(`.article-content`), {
            delimiters: [
                { left: "$$", right: "$$", display: true },
                { left: "$", right: "$", display: false },
                { left: "\\(", right: "\\)", display: false },
                { left: "\\[", right: "\\]", display: true }
            ]
        });})
</script>
    
</article>

    

    

     
    
        
    

    <footer class="site-footer">
    <section class="copyright">
        &copy; 
        
            2018 - 
        
        2024 Charent的博客
    </section>
    
    <section class="powerby">
         <br />
        
    </section>
</footer>


    
<div class="pswp" tabindex="-1" role="dialog" aria-hidden="true">

    
    <div class="pswp__bg"></div>

    
    <div class="pswp__scroll-wrap">

        
        <div class="pswp__container">
            <div class="pswp__item"></div>
            <div class="pswp__item"></div>
            <div class="pswp__item"></div>
        </div>

        
        <div class="pswp__ui pswp__ui--hidden">

            <div class="pswp__top-bar">

                

                <div class="pswp__counter"></div>

                <button class="pswp__button pswp__button--close" title="Close (Esc)"></button>

                <button class="pswp__button pswp__button--share" title="Share"></button>

                <button class="pswp__button pswp__button--fs" title="Toggle fullscreen"></button>

                <button class="pswp__button pswp__button--zoom" title="Zoom in/out"></button>

                
                
                <div class="pswp__preloader">
                    <div class="pswp__preloader__icn">
                        <div class="pswp__preloader__cut">
                            <div class="pswp__preloader__donut"></div>
                        </div>
                    </div>
                </div>
            </div>

            <div class="pswp__share-modal pswp__share-modal--hidden pswp__single-tap">
                <div class="pswp__share-tooltip"></div>
            </div>

            <button class="pswp__button pswp__button--arrow--left" title="Previous (arrow left)">
            </button>

            <button class="pswp__button pswp__button--arrow--right" title="Next (arrow right)">
            </button>

            <div class="pswp__caption">
                <div class="pswp__caption__center"></div>
            </div>

        </div>

    </div>

</div><script 
                src="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/photoswipe.min.js"integrity="sha256-ePwmChbbvXbsO02lbM3HoHbSHTHFAeChekF1xKJdleo="crossorigin="anonymous"
                defer
                >
            </script><script 
                src="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/photoswipe-ui-default.min.js"integrity="sha256-UKkzOn/w1mBxRmLLGrSeyB4e1xbrp4xylgAWb3M42pU="crossorigin="anonymous"
                defer
                >
            </script><link 
                rel="stylesheet" 
                href="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/default-skin/default-skin.css"integrity="sha256-c0uckgykQ9v5k&#43;IqViZOZKc47Jn7KQil4/MP3ySA3F8="crossorigin="anonymous"
            ><link 
                rel="stylesheet" 
                href="https://cdn.jsdelivr.net/npm/photoswipe@4.1.3/dist/photoswipe.css"integrity="sha256-SBLU4vv6CA6lHsZ1XyTdhyjJxCjPif/TRkjnsyGAGnE="crossorigin="anonymous"
            >

            </main>
    <aside class="sidebar right-sidebar sticky">
        
            
                
    <section class="widget archives">
        <div class="widget-icon">
            <svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-hash" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
  <path stroke="none" d="M0 0h24v24H0z"/>
  <line x1="5" y1="9" x2="19" y2="9" />
  <line x1="5" y1="15" x2="19" y2="15" />
  <line x1="11" y1="4" x2="7" y2="20" />
  <line x1="17" y1="4" x2="13" y2="20" />
</svg>



        </div>
        <h2 class="widget-title section-title">目录</h2>
        
        <div class="widget--toc">
            <nav id="TableOfContents">
  <ol>
    <li><a href="#1-将layer_norm更换为rms_norm">1. 将Layer_norm更换为RMS_norm</a></li>
    <li><a href="#2-q在与k相乘之前先使用rope进行位置编码">2. Q在与K相乘之前，先使用RoPE进行位置编码</a></li>
    <li><a href="#3-引入kv-cache并采用group-query-attention">3. 引入KV Cache，并采用Group Query Attention</a></li>
    <li><a href="#4-用silu激活函数代替relugelu">4. 用SiLU激活函数代替RELU/GELU</a></li>
    <li><a href="#5-使用分组查询注意力-grouped-query-attention-gqa">5. 使用分组查询注意力 (Grouped Query Attention, GQA)</a></li>
  </ol>
</nav>
        </div>
    </section>

            
        
    </aside>


        </div>
        <script 
                src="https://cdn.jsdelivr.net/npm/node-vibrant@3.1.5/dist/vibrant.min.js"integrity="sha256-5NovOZc4iwiAWTYIFiIM7DxKUXKWvpVEuMEPLzcm5/g="crossorigin="anonymous"
                
                >
            </script><script type="text/javascript" src="/ts/main.js" defer></script>
<script>
    (function () {
        const customFont = document.createElement('link');
        customFont.href = "https://fonts.googleapis.com/css2?family=Lato:wght@300;400;700&display=swap";

        customFont.type = "text/css";
        customFont.rel = "stylesheet";

        document.head.appendChild(customFont);
    }());
</script>

    </body>
</html>
